#ifndef LQR_TEST__SIMPLE_LQR_SOLVER_HPP_
#define LQR_TEST__SIMPLE_LQR_SOLVER_HPP_
#include <iostream>
#include <vector>
#include <math.h>
#include <stdint.h>
#include <string>
#include <algorithm>

#include <Eigen/Dense>
#include <Eigen/Core>
#include <Eigen/QR>

#include "agv_model.hpp"
#include "angles/angles.h"

typedef struct Output
{
    double linear;
    double angular;
};

typedef Eigen::Matrix<double, 2, 1> Matrix2x1;
typedef Eigen::Matrix<double, 2, 2> Matrix2x2;
typedef Eigen::Matrix<double, 2, 3> Matrix2x3;
typedef Eigen::Matrix<double, 3, 1> Matrix3x1;
typedef Eigen::Matrix<double, 3, 2> Matrix3x2;
typedef Eigen::Matrix<double, 3, 3> Matrix3x3;
typedef Eigen::Matrix<double, 3, 6> Matrix3x6;

class SimpleLQRSolver
{
private:
    double tolerance_;
    int iter_max_ = 100;
    double length_;
    double dt_;
    Matrix3x3 Q_;
    Matrix2x2 R_;

    Matrix3x3 A_;
    Matrix3x2 B_;
    Matrix3x3 P_;

    RobotState robot_;
    WayPoint ref_;

public:
    SimpleLQRSolver(double length, double freq)
    {
        tolerance_ = 1e-2;
        length_ = length;
        dt_ = 1 / freq;

        Q_ << 30.0, 0.0, 0.0,
            0.0, 30.0, 0.0,
            0.0, 0.0, 5.0;
        R_ << 10.0, 0.0,
            0.0, 1.0;
    }

    ~SimpleLQRSolver()
    {
    }

    void initial(RobotState cur_pose, WayPoint ref_pose, Output ref_u)
    {
        A_ = Matrix3x3::Identity();
        A_(0, 2) = -dt_ * ref_u.linear * sin(ref_pose.yaw);
        A_(1, 2) = dt_ * ref_u.linear * cos(ref_pose.yaw);

        B_ = Matrix3x2::Zero();
        B_(0, 0) = dt_ * cos(ref_pose.yaw);
        B_(1, 0) = dt_ * sin(ref_pose.yaw);
        B_(2, 0) = 0.0;
        B_(2, 1) = dt_;

        robot_ = cur_pose;
        ref_ = ref_pose;
    }

    Matrix3x1 computeError()
    {
        Matrix3x1 e;
        e(0) = robot_.x - ref_.x;
        e(1) = robot_.y - ref_.y;
        e(2) = angles::shortest_angular_distance(ref_.yaw, robot_.yaw);
        return e;
    }

    bool isControllable()
    {
        // 构建可控性矩阵 [B, A*B, A^2*B]
        Matrix3x6 ctrb;
        Matrix3x2 AB = B_;
        ctrb.block<3, 2>(0, 0) = B_;
        AB = A_ * AB;
        ctrb.block<3, 2>(0, 2) = AB;
        AB = A_ * AB;
        ctrb.block<3, 2>(0, 4) = AB;

        Eigen::FullPivLU<Matrix3x6> lu(ctrb);
        int rank = lu.rank();
        return (rank == 3);
    }

    bool solveRiccati()
    {
        // 先检查可控性
        if (!isControllable())
        {
            std::cerr << "[LQR] Warning: (A, B) is not fully controllable. rank < 3\n";
            return false;
        }

        P_ = Q_; // initialize

        const double eps_reg = 1e-6; // 基本正则项

        Matrix3x3 AdT = A_.transpose();
        Matrix2x3 BdT = B_.transpose();
        double diff;
        for (int i = 0; i < iter_max_; ++i)
        {
            // Build M = R + B^T P B, 若病态则做正则化
            Matrix2x2 M = R_ + BdT * P_ * B_;
            // 保证 M 对称
            M = 0.5 * (M + M.transpose());
            // 检查最小特征值（而不是仅仅 determinant），更稳健
            Eigen::SelfAdjointEigenSolver<Matrix2x2> es(M);
            double min_eig = es.eigenvalues().minCoeff();
            if (min_eig < eps_reg)
            {
                // 增加对角正则项，防止奇异
                double add = std::max(eps_reg, -min_eig + eps_reg);
                M.diagonal().array() += add;
                // 也同时在 R_ 上反映（不永久改变 R_，仅本次计算）
            }
            // -- discrete solver --
            // 计算逆：优先使用稳定的解法（CHOLESKY if PD else pseudoInverse）
            Matrix2x2 M_inv;
            // 使用 Eigen 的 LLT 判断是否正定
            Eigen::LLT<Matrix2x2> llt(M);
            if (llt.info() == Eigen::Success)
            {
                M_inv = llt.solve(Matrix2x2::Identity());
            }
            else
            {
                // 退化：使用伪逆（更稳健）
                Eigen::CompleteOrthogonalDecomposition<Matrix2x2> cod(M);
                M_inv = cod.pseudoInverse();
            }

            Matrix3x3 P_next = AdT * P_ * A_ - AdT * P_ * B_ * M_inv * BdT * P_ * A_ + Q_;
            // 强制对称（数值误差）
            P_next = 0.5 * (P_next + P_next.transpose());

            diff = (P_next - P_).cwiseAbs().maxCoeff();
            P_ = P_next;
            // 防止发散（P 非法）
            Eigen::SelfAdjointEigenSolver<Matrix3x3> esP(P_);
            if (esP.eigenvalues().minCoeff() < 0.0)
            {
                std::cerr << "[LQR] P became indefinite, aborting Riccati\n";
                return false;
            }
            if (diff < tolerance_)
            {
                return true;
            }
        }
        // 超过最大迭代次数仍未收敛
        std::cerr << "[LQR] Riccati did not converge within iter_max_\n";
        std::cerr << "diff" << diff << std::endl;
        return false;
    }

    Output calcu()
    {
        // 计算代价函数矩阵 P_
        bool success = solveRiccati();
        if (!success)
        {
            std::cerr << "solve riccati failed" << std::endl;
            // std::cerr << "A: \n" << A_ << std::endl;
            // std::cerr << "B: \n" << B_ << std::endl;
        }

        // 计算状态误差 x
        Matrix3x1 e = computeError();

        // ----控制律：u = -K * e------
        // 构造控制增益 K：如果 Riccati 成功使用 P_；若失败，采用稳健 fallback（加大 R 并用伪逆）
        Matrix2x3 K;
        {
            Matrix2x3 BdT = B_.transpose();
            Matrix2x2 temp = R_ + BdT * P_ * B_;
            // 强制对称
            temp = 0.5 * (temp + temp.transpose());

            // 若 temp 条件差或接近奇异，增加正则化到对角
            Eigen::SelfAdjointEigenSolver<Matrix2x2> es(temp);
            double min_eig = es.eigenvalues().minCoeff();
            if (min_eig < 1e-6)
            {
                // 在无法求解 Riccati 或 temp 数值差时，增加 R 的对角（局部修改）
                temp.diagonal().array() += 1e-3;
            }

            // 使用伪逆以保证健壮性
            Eigen::CompleteOrthogonalDecomposition<Matrix2x2> cod(temp);
            Matrix2x2 temp_inv = cod.pseudoInverse();

            K = temp_inv * BdT * P_ * A_;
        }
        Matrix2x1 u = -K * e;

        Output cmd;
        cmd.linear = u(0);
        cmd.angular = u(1);
        return cmd;
    }
};
#endif
